package com.fy.fyspace.common.decorator;

import com.fy.fyspace.common.limiter.RateLimiter;
import com.fy.fyspace.model.entity.FileInfo;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.InputStream;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.TimeUnit;

/**
 * 限流下载装饰器
 */
@Slf4j
public class RateLimitedDownloadDecorator {
    private static final int BUFFER_SIZE = 8192;
    private final RateLimiter rateLimiter;

    public RateLimitedDownloadDecorator(RateLimiter rateLimiter) {
        this.rateLimiter = rateLimiter;
    }

    public void downloadWithRateLimit(InputStream inputStream, HttpServletResponse response, FileInfo fileInfo) throws IOException {
        response.setContentType("application/octet-stream");
        response.setHeader(HttpHeaders.CONTENT_DISPOSITION,
                "attachment;filename=" + URLEncoder.encode(fileInfo.getFileName(), StandardCharsets.UTF_8.name()));

        try (ServletOutputStream outputStream = response.getOutputStream()) {
            byte[] buffer = new byte[BUFFER_SIZE];
            int bytesRead;
            
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                if (bytesRead > 0) {
                    // 等待获取令牌后再写入数据
                    rateLimiter.tryAcquire(bytesRead, 1, TimeUnit.SECONDS);
                    outputStream.write(buffer, 0, bytesRead);
                    outputStream.flush();
                }
            }
        } catch (Exception e) {
            log.error("File download failed for file: {}", fileInfo.getFileName(), e);
            throw new IOException("Download failed", e);
        } finally {
            if (inputStream != null) {
                try {
                    inputStream.close();
                } catch (IOException e) {
                    log.error("Error closing input stream", e);
                }
            }
        }
    }
} 